#!/opt/anaconda3/bin/python
import sys, os
import pyfaidx


def cal_Nx0(lengthList, Nlength):
    Sum = 0
    for i in lengthList:
        Sum = Sum + i
        if Sum > Nlength:
            return i
        else:
            continue


def statFa(fasta):
    fa = pyfaidx.Faidx(fasta)
    fain = pyfaidx.Fasta(fasta)
    lengthList = [fa.index[record].rlen for record in fain.keys()]
    lengthList.sort(reverse=True)
    # print(lengthList)
    totalBase = sum(lengthList)
    numOfSeq = len(lengthList)
    N500bp = 0
    N2000bp = 0
    N90, N80, N70, N60, N50 = [0] * 5
    maxLength = 0
    minLength = 10000000000000
    meanLength = totalBase / numOfSeq
    for num in lengthList:
        if num > 500:
            N500bp += 1
        if num > 2000:
            N2000bp += 1
        if num > maxLength:
            maxLength = num
        if num < minLength:
            minLength = num
    N50 = cal_Nx0(lengthList, totalBase * 0.5)
    N60 = cal_Nx0(lengthList, totalBase * 0.6)
    N70 = cal_Nx0(lengthList, totalBase * 0.7)
    N80 = cal_Nx0(lengthList, totalBase * 0.8)
    N90 = cal_Nx0(lengthList, totalBase * 0.9)

    print(f"number of seq:  {numOfSeq}")
    print(f"min length: {minLength}")
    print(f"max length: {maxLength}")
    print(f"Total Size: {totalBase}")
    print(f"N90: {N90}")
    print(f"N80: {N80}")
    print(f"N70: {N70}")
    print(f"N60: {N60}")
    print(f"N50: {N50}")
    print(f"avg length: {int(meanLength)}")
    print(f"Total Num(>500bp): {N500bp}")
    print(f"Total Num(>2kb): {N2000bp}")


if __name__ == "__main__":
    try:
        statFa(sys.argv[1])
    except:
        print("Usage:FaStat fastaFile")
